execfile('v4/base.py')

class RaiseF(Base):
    def __init__(self):
        self.bar_count_default = 50
        self.end_date_str = '20250831'  # 如果需要可以取消注释
        # self.start_date_str = '2025-06-16'  # 如果需要可以取消注释
    
    '''获取涨幅特征'''
    def mark_raise(self,df):
        raise_df = df.copy()
        # 预排序避免重复排序开销
        raise_df = raise_df.sort_values(['code', 'date'])
        raise_df['date'] = pd.to_datetime(raise_df['date'])
        
        # 1、计算1,2，5,10,20,40,60日涨幅特征；
        close_groups = raise_df.groupby('code')['close']
        periods = [1,2,5,10,20,40,60]
        for p in periods:
            raise_df[f'r{p}'] = close_groups.pct_change(p).mul(100)
        
        def _cal_gt_flags(g):
            current_code = g.name
            processed = g.assign(
                code=current_code,  # 重建code列
                gt4=lambda x: x['quote_rate'].rolling(10, min_periods=1).max() > 4,
                gt7=lambda x: x['quote_rate'].rolling(10, min_periods=1).max() > 7,
                gt9=lambda x: x['quote_rate'].rolling(10, min_periods=1).max() > 9
            )
            return processed
    
        # 2、生成特征gt4，gt7、gt9。（10日内是否出现过大于4,7,9）
        raise_df = (
            raise_df.groupby('code', group_keys=True)  # 保留分组键元信息
            .apply(_cal_gt_flags, include_groups=False)  # 显式排除分组列
            .reset_index(drop=True)  # 清除残留的层级索引
        )
        # 3、涨停相关计算 (当天)
        raise_df['is_zt'] = (raise_df['close'] == raise_df['high_limit']).astype('bool')
        raise_df['is_gt4'] = (raise_df['quote_rate'] > 4).astype('bool')
        raise_df['is_gt7'] = (raise_df['quote_rate'] > 7).astype('bool')
        raise_df['is_gt9'] = (raise_df['quote_rate'] > 9).astype('bool')

        raise_df['gt4'] = raise_df.groupby('code')['is_gt4'].transform(
            lambda s: s.rolling(10, min_periods=1).sum().astype('int64')
        )
        raise_df['gt7'] = raise_df.groupby('code')['is_gt7'].transform(
            lambda s: s.rolling(10, min_periods=1).sum().astype('int64')
        )
        raise_df['gt9'] = raise_df.groupby('code')['is_gt9'].transform(
            lambda s: s.rolling(10, min_periods=1).sum().astype('int64')
        )
        
        # 创建辅助列：标记涨停序列的开始
        def calculate_consecutive_zt_efficient(df):
            """
            高效迭代计算连续涨停天数（使用分组应用）
            """
            # 确保数据按股票代码和日期排序
            df = df.sort_values(['code', 'date']).copy()
            
            # 使用分组应用计算连续涨停天数
            def calc_for_one_stock(group):
                group = group.sort_values('date')
                consecutive = 0
                result = []
                
                for i in range(len(group)):
                    if group.iloc[i]['is_zt']:
                        if i == 0 or not group.iloc[i-1]['is_zt']:
                            consecutive = 1
                        else:
                            consecutive += 1
                    else:
                        consecutive = 0
                    result.append(consecutive)
                
                return pd.Series(result, index=group.index)
            
            # 对每个股票分组应用计算函数
            df['zt_days'] = df.groupby('code', group_keys=False)[df.columns].apply(calc_for_one_stock)

            return df

        raise_df = calculate_consecutive_zt_efficient(raise_df)

        raise_df['is_zt_start'] = raise_df['is_zt'] & ~raise_df.groupby('code')['is_zt'].shift(1).fillna(False)
        # raise_df['zt_days'] = raise_df.groupby('code')['is_zt_start'].cumsum()
        # raise_df['zt_days'] = raise_df.groupby(['code', 'zt_days']).cumcount() + 1

    

        raise_df['zt10'] = raise_df.groupby('code')['is_zt'].transform(
            lambda s: s.rolling(10, min_periods=1).sum().astype('int64')
        )
        raise_df['zt20'] = raise_df.groupby('code')['is_zt'].transform(
        lambda s: s.rolling(20, min_periods=1).sum().astype('int64')
        )

        # 4、跳空高开特征
        raise_df['prev_close'] = raise_df.groupby('code')['close'].shift(1)
        raise_df['open_ratio'] = (raise_df['open'] - raise_df['prev_close']) / raise_df['prev_close'] * 100
       
        # 5、昨日交易特征
        def _cal_yesterday_feature(df):
            # 按股票分组并确保日期升序排列
            df = df.sort_values(by=['code', 'date'])
            # 计算昨日涨幅：将r1滞后1天
            df['pre_r1'] = df.groupby('code')['r1'].shift(1)
            df['pre_t_rate'] = df.groupby('code')['t_rate'].shift(1) #昨日换手率
    
            return df
        raise_df = raise_df.groupby('code', group_keys=False)[raise_df.columns].apply(_cal_yesterday_feature)
    
        
        # 6、周、月涨幅
        def _cal_month_raise(g):
            g['month'] = g['date'].dt.to_period('M')
            first_close = g.groupby('month')['close'].transform('first')
            first_quote = g.groupby('month')['quote_rate'].transform('first')
            g['m_raise'] = (g['close'] - first_close)/first_close *100 + first_quote
            return g.drop(columns='month')
        # 保留其他必要计算（月涨幅、涨停数等）
        raise_df = raise_df.groupby('code', group_keys=False)[raise_df.columns].apply(_cal_month_raise)
        ''' 计算周涨幅 '''
        def calculate_weekly_raise2(df):
            # 预处理：确保按日期排序
            df = df.sort_values(['code', 'date']).copy()
            df['date'] = pd.to_datetime(df['date'])
            
            def _cal_week(g):
                # 生成精准周标记（ISO周）
                g['week'] = g['date'].dt.isocalendar().year.astype(str) + '-' + \
                           g['date'].dt.isocalendar().week.astype(str).str.zfill(2)
                
                # 周内按日期排序后累计
                g = g.sort_values('date')
                g['w_raise'] = g.groupby('week')['quote_rate'].cumsum()
                return g.drop(columns='week')
            
            return df.groupby('code', group_keys=False)[df.columns].apply(_cal_week)


        raise_df =  calculate_weekly_raise2(raise_df)

        # 7、构建四级信号体系（条件优先级从高到低）
        conditions = [
            raise_df['quote_rate'] >= 9,                        # 条件1：涨幅≥9%
            raise_df['quote_rate'].between(4, 9, inclusive="left"),  # 条件2：4% ≤ 涨幅 <9%
            (raise_df['quote_rate'] < 4) & (raise_df['gt4'] >0)  # 条件3：<4%但历史存在>4%
        ]
        choices = [1, 2, 3]
        raise_df['raise_buy'] = np.select(conditions, choices, default=4)

    
        raise_df.drop(['prev_close','prev_high'], axis=1,errors='ignore', inplace=True)
        raise_df = self.drop_base_columns(raise_df)
        return raise_df.round(2)

   
    
